package com.yangye.niodemo.multiplexing.threadgroup;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.Channel;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.Iterator;
import java.util.Set;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * @author yangye
 * @date 2020/12/24
 */
public class SelectorThread implements Runnable {

	public SelectorThreadGroup stg;

	public Selector selector;

	public LinkedBlockingQueue<Channel> queue = new LinkedBlockingQueue<>();

	public SelectorThread(SelectorThreadGroup stg) {
		this.stg = stg;
		try {
			selector = Selector.open();
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	@Override
	public void run() {
		while (true) {
			try {
				if (selector.select() > 0) {
					Set<SelectionKey> selectionKeys = selector.selectedKeys();
					Iterator<SelectionKey> iter = selectionKeys.iterator();
					while (iter.hasNext()) {
						SelectionKey key = iter.next();
						iter.remove();
						if (key.isAcceptable()) {// 客户端连接事件
							handleAccept(key);
						} else if (key.isReadable()) {// 接收到消息
							handleRead(key);
						}
					}
				}

				// 将新增的channel注册到多路复用器上
				if (!queue.isEmpty()) {
					Channel c = queue.take();
					if (c instanceof ServerSocketChannel) {
						ServerSocketChannel server = (ServerSocketChannel) c;
						server.register(selector, SelectionKey.OP_ACCEPT);
						System.out.println(Thread.currentThread().getName() + " register listen");
					} else if (c instanceof SocketChannel) {
						SocketChannel client = (SocketChannel) c;
						ByteBuffer buffer = ByteBuffer.allocateDirect(4096);
						client.register(selector, SelectionKey.OP_READ, buffer);
						System.out.println(Thread.currentThread().getName() + " register client: " + client.getRemoteAddress());
					}
				}
			} catch (IOException | InterruptedException e) {
				e.printStackTrace();
			}
		}
	}

	private void handleRead(SelectionKey key) {
		SocketChannel client = (SocketChannel) key.channel();
		ByteBuffer buffer = (ByteBuffer) key.attachment();
		buffer.clear();

		int read = 0;
		try {
			while (true) {
				read = client.read(buffer);
				if (read > 0) {
					buffer.flip();
					while (buffer.hasRemaining()) {
						client.write(buffer);
					}
					buffer.clear();
				} else if (read == 0) {
					break;
				} else {
					client.close();
					break;
				}
			}
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	private void handleAccept(SelectionKey key) {
		ServerSocketChannel server = (ServerSocketChannel) key.channel();
		try {
			SocketChannel client = server.accept();
			client.configureBlocking(false);

			// 注册到worker线程组
			stg.registerWorker(client);

		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	public static void main(String[] args) {
		System.out.println(Runtime.getRuntime().availableProcessors());
	}
}
